import numpy as np
from matplotlib import pyplot as plt
from skimage import io
from skimage.color import rgb2lab, lab2rgb
from sklearn.neighbors import KNeighborsRegressor

block_size = 1
# 训练集：数据（3×3的小方格），标签（中心点的A，B）
def read_style_image(file_name,block_size=1):
    img = io.imread(file_name)
    img = rgb2lab(img)
    w,h = img.shape[:2]
    X = []
    Y = []
    for x in range(block_size,w-block_size):
        for y in range(block_size,h-block_size):
            X.append(img[x-block_size:x+block_size+1,y-block_size:y+block_size+1,0].flatten())
            Y.append(img[x,y,1:])
    return X,Y

file_name = '../data/style_transfer/vangogh/00001.jpg'
X,Y = read_style_image(file_name,block_size)

# 训练KNN
knn = KNeighborsRegressor(n_neighbors=4,weights='distance')
knn.fit(X,Y)

def rebuild(img,block_size=1):
    img = rgb2lab(img)
    w,h = img.shape[:2]
    X = []
    for x in range(block_size,w-block_size):
        for y in range(block_size,h-block_size):
            window = img[x-block_size:x+block_size+1,y-block_size:y+block_size+1,0].flatten()
            X.append(window)
    X = np.array(X)

    pred_ab = knn.predict(X)

    photo = np.zeros([w,h,3])
    photo[:,:,0] = img[:,:,0]
    photo[block_size:w-block_size,block_size:h-block_size,1:] = pred_ab.reshape(w-2*block_size,h-2*block_size,2)
    return photo

input_photo = "../data/style_transfer/input.jpg"
input_img = io.imread(input_photo)
output_photo = rebuild(input_img)
output_photo = lab2rgb(output_photo)
plt.imshow(output_photo)
plt.show()

